import argparse
import logging
import os
import matplotlib
import matplotlib.pyplot as plt 
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms

from unet3d.model import UNet2D, UNet3D
from dataset import BasicDataset
from dice_loss import dice_coeff

img_size = (256, 256)
out_channel = 1

def predict_img(net,
                full_img,
                device,
                img_size=img_size,
                out_threshold=0.5):
    net.eval()

    img = torch.from_numpy(BasicDataset.preprocess(full_img, img_size=img_size))

    img = img.unsqueeze(0)
    img = img.to(device=device, dtype=torch.float32)

    with torch.no_grad():
        output = net(img)
        output = output.squeeze(2)

        if out_channel > 1:
            probs = F.softmax(output, dim=1)
        else:
            probs = torch.sigmoid(output)

        probs = probs.squeeze(0)
        print(probs.shape)
        
        tf = transforms.Compose(
            [
                transforms.ToPILImage(),
                transforms.Resize(full_img.size[1]),
                transforms.ToTensor()
            ]
        )

        probs = tf(probs.cpu())
        full_mask = probs.squeeze().cpu().numpy()

    return full_mask > out_threshold


def mask_to_image(mask):
    return Image.fromarray((mask * 255).astype(np.uint8))

if __name__ == "__main__":
   
    file_path = "./DSB_data/1f0008060150b5b93084ae2e4dabd160ab80a95ce8071a321b80ec4e33b58aca/images/1f0008060150b5b93084ae2e4dabd160ab80a95ce8071a321b80ec4e33b58aca.png"
    net = UNet2D(in_channels=1, out_channels=1)


    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    logging.info(f'Using device {device}')
    net.to(device=device)
    net.load_state_dict(torch.load("./unet2D_epoch11.pth", map_location=device))

    logging.info("Model loaded !")

    img = Image.open(file_path).convert('L')

    img_tensor = BasicDataset.preprocess(img, img_size=(256, 256))
    print(img_tensor.shape)
    img_tensor = torch.from_numpy(img_tensor).float()
    img_tensor = img_tensor.unsqueeze(0)

    pred_mask = net(img_tensor)
    pred_mask = pred_mask.squeeze(2)

    print(pred_mask.shape)
    
    # print(torch.sigmoid(pred_mask[0, 0, :10, :40]))

    pred_mask = (torch.sigmoid(pred_mask) > 0.5).float()
    print(pred_mask[0, 0, :10, :40])

    file_path = "./DSB_data/1f0008060150b5b93084ae2e4dabd160ab80a95ce8071a321b80ec4e33b58aca/save_masks/1f0008060150b5b93084ae2e4dabd160ab80a95ce8071a321b80ec4e33b58aca.png"
    true_mask = Image.open(file_path).convert("L")
    true_mask_tensor = BasicDataset.preprocess(true_mask, img_size=(256, 256))
    print(true_mask_tensor.shape)
    true_mask_tensor = torch.from_numpy(true_mask_tensor).float()
    
    # print(true_mask_tensor.shape)
    true_mask_tensor = (true_mask_tensor > 0.5).float()

    print(true_mask_tensor[0, 0, :10, :40])
    print(dice_coeff(pred_mask, true_mask_tensor))

    mask = predict_img(net=net,
                        full_img=img,
                        img_size=img_size,
                        device=device)
    print(mask.shape)

    im = Image.fromarray(mask)
    im.show()

  
